package com.tsyz.servlet.util;

import java.util.UUID;
import java.util.concurrent.Callable;

/**
 * TraceId工具类，用于生成和管理请求traceId
 */
public class TraceIdUtil {

    public static final String TRACE_ID = "traceId";

    private static final ThreadLocal<String> TRACE_ID_LOCAL = new ThreadLocal<>();

    /**
     * 生成 TraceID
     *
     * @return UUID格式的TraceId
     */
    public static String generateTraceId() {
        return UUID.randomUUID().toString().replace("-", "").toUpperCase();
    }

    /**
     * 初始化一个新的TraceId并设置到当前线程
     * 这是系统中唯一应该直接生成TraceId的地方
     *
     * @return 新生成的TraceId
     */
    public static String initTraceId() {
        // 清理可能存在的旧TraceId
        clear();

        // 生成新的TraceId
        String traceId = generateTraceId();
        TRACE_ID_LOCAL.set(traceId);

        return traceId;
    }

    /**
     * 返回当前线程对应的traceId
     *
     * @return traceId字符串
     */
    public static String getTraceId() {
        return TRACE_ID_LOCAL.get();
    }

    /**
     * 设置当前线程的traceId
     *
     * @param traceId traceId字符串
     */
    public static void setTraceId(String traceId) {
        TRACE_ID_LOCAL.set(traceId);
    }

    /**
     * 清理当前线程的traceId
     */
    public static void clear() {
        TRACE_ID_LOCAL.remove();
    }

    /**
     * 在新的线程中继承父线程的TraceId并执行任务
     *
     * @param runnable 要执行的任务
     */
    public static void executeWithTraceId(Runnable runnable) {
        String traceId = getTraceId();
        Runnable wrappedRunnable = () -> {
            try {
                setTraceId(traceId);
                runnable.run();
            } finally {
                clear();
            }
        };
        // 直接在当前线程执行，也可以提交到线程池
        wrappedRunnable.run();
    }

    /**
     * 在新的线程中继承父线程的TraceId并执行任务，支持返回值
     *
     * @param callable 要执行的任务
     * @param <T>      返回值类型
     * @return callable的返回值
     */
    public static <T> T callWithTraceId(Callable<T> callable) throws Exception {
        String traceId = getTraceId();
        Callable<T> wrappedCallable = () -> {
            try {
                setTraceId(traceId);
                return callable.call();
            } finally {
                clear();
            }
        };
        return wrappedCallable.call();
    }

    /**
     * 包装Runnable任务，使其在执行时继承父线程的TraceId
     *
     * @param runnable 原始任务
     * @return 包装后的任务
     */
    public static Runnable wrap(Runnable runnable) {
        String traceId = getTraceId();
        return () -> {
            try {
                setTraceId(traceId);
                runnable.run();
            } finally {
                clear();
            }
        };
    }

    /**
     * 包装Callable任务，使其在执行时继承父线程的TraceId
     *
     * @param callable 原始任务
     * @param <T>      返回值类型
     * @return 包装后的任务
     */
    public static <T> Callable<T> wrap(Callable<T> callable) {
        String traceId = getTraceId();
        return () -> {
            try {
                setTraceId(traceId);
                return callable.call();
            } finally {
                clear();
            }
        };
    }
}
